//
// Created by clarence on 2021/4/14.
//

#include "ros/ros.h"
#include "vied_map_array_improve_5_cluster_velocity.h"
#include <pcl/common/transforms.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <sensor_msgs/PointCloud2.h>
#include "geometry_msgs/Pose.h"
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>
#include <geometry_msgs/TwistStamped.h>
#include "gazebo_msgs/ModelStates.h"
#include "geometry_msgs/PoseStamped.h"
#include <queue>
#include "nav_msgs/Odometry.h"


VideMap my_map;

queue<double> pose_att_time_queue;
queue<Eigen::Vector3f> uav_position_global_queue;
queue<Eigen::Quaternionf> uav_att_global_queue;

const unsigned int MAX_POINT_NUM  = 5000;
float point_clouds[MAX_POINT_NUM*3];

float x_min = -MAP_LENGTH_VOXEL_NUM * VOXEL_RESOLUTION / 2;
float x_max = MAP_LENGTH_VOXEL_NUM * VOXEL_RESOLUTION / 2;
float y_min = -MAP_WIDTH_VOXEL_NUM * VOXEL_RESOLUTION / 2;
float y_max = MAP_WIDTH_VOXEL_NUM * VOXEL_RESOLUTION / 2;
float z_min = -MAP_HEIGHT_VOXEL_NUM * VOXEL_RESOLUTION / 2;
float z_max = MAP_HEIGHT_VOXEL_NUM * VOXEL_RESOLUTION / 2;

ros::Publisher cloud_pub, map_center_pub, gazebo_model_states_pub, current_velocity_pub, single_object_velocity_pub, single_object_velocity_truth_pub;
ros::Publisher future_status_pub;
gazebo_msgs::ModelStates ground_truth_model_states;
int ground_truth_updated = 0;
float res = 0.2;  // Smaller res will get better tracking result but is slow.

int inRange(float &low, float &high, float &x)
{
    if(x > low && x < high){
        return 1;
    }else{
        return 0;
    }
}

void colorAssign(int &r, int &g, int &b, float v, float value_min=0.f, float value_max=1.f, int reverse_color=0)
{
    v = std::max(v, value_min);
    v = std::min(v, value_max);

    float v_range = value_max - value_min;
    int value = floor((v - value_min) / v_range * 240); // Mapping 0~1.0 to 0~240
    value = std::min(value, 240);

    if(reverse_color){
        value = 240 - value;
    }

    int section = value / 60;
    float float_key = (value % 60) / (float)60 * 255;
    int key = floor(float_key);
    int nkey = 255 - key;

    switch(section) {
        case 0: // G increase
            r = 255;
            g = key;
            b = 0;
            break;
        case 1: // R decrease
            r = nkey;
            g = 255;
            b = 0;
            break;
        case 2: // B increase
            r = 0;
            g = 255;
            b = key;
            break;
        case 3: // G decrease
            r = 0;
            g = nkey;
            b = 255;
            break;
        case 4:
            r = 0;
            g = 0;
            b = 255;
            break;
        default: // White
            r = 255;
            g = 255;
            b = 255;
    }

}


pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_filtered(new pcl::PointCloud<pcl::PointXYZ>());
void cloudCallback(const sensor_msgs::PointCloud2ConstPtr& cloud)
{

    /// Simple synchronizer
    Eigen::Vector3f uav_position;
    Eigen::Quaternionf uav_att;

    while(!pose_att_time_queue.empty()){   //Sychronize pose by queue
        double time_stamp_pose = pose_att_time_queue.front();
        if(time_stamp_pose >= cloud->header.stamp.toSec() && time_stamp_pose <= cloud->header.stamp.toSec() + 1.0){
            uav_att = uav_att_global_queue.front();
            uav_position = uav_position_global_queue.front();
            ROS_INFO_THROTTLE(3.0, "cloud mismatch time = %lf", cloud->header.stamp.toSec() - time_stamp_pose);
            break;
        }
        pose_att_time_queue.pop();
        uav_att_global_queue.pop();
        uav_position_global_queue.pop();
    }
    // Record gazebo models' states for evaluation
    gazebo_msgs::ModelStates ground_truth_model_states_this = ground_truth_model_states;

    /// Point cloud process
    static double last_time = ros::Time::now().toSec();
    double this_time = ros::Time::now().toSec();

    // convert cloud to pcl form
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_in(new pcl::PointCloud<pcl::PointXYZ>());
    pcl::fromROSMsg(*cloud, *cloud_in);

    // down-sample for all
    pcl::VoxelGrid<pcl::PointXYZ> sor;
    sor.setInputCloud(cloud_in);
    sor.setLeafSize(res, res, res);
    sor.filter(*cloud_filtered);

    int useful_point_num = 0;
    for(int i=0; i<cloud_filtered->width; i++){
        float x = cloud_filtered->points.at(i).z;
        float y = -cloud_filtered->points.at(i).x;
        float z = -cloud_filtered->points.at(i).y;

        if(inRange(x_min, x_max, x) && inRange(y_min, y_max, y) && inRange(z_min, z_max, z))
        {
            point_clouds[useful_point_num*3] = x;
            point_clouds[useful_point_num*3+1] = y;
            point_clouds[useful_point_num*3+2] = z;
            ++ useful_point_num;

            if(useful_point_num >= MAX_POINT_NUM){
                break;
            }
        }
    }


    /// Update
    clock_t start1, finish1;
    start1 = clock();

    my_map.update(useful_point_num, 3, point_clouds,
                  uav_position.x(), uav_position.y(), uav_position.z(), (float)(this_time-last_time),
                  uav_att.w(), uav_att.x(), uav_att.y(), uav_att.z());

    finish1 = clock();
    double duration1 = (double)(finish1 - start1) / CLOCKS_PER_SEC;
    printf( "****** Map update time %f seconds\n", duration1);


    /// Get occupancy status
    clock_t start2, finish2;
    start2 = clock();

    int occupied_num=0;
    pcl::PointCloud<pcl::PointXYZ> cloud_to_publish;
    sensor_msgs::PointCloud2 cloud_to_pub_transformed;
    static float future_status[VOXEL_NUM][PREDICTION_TIMES];

    my_map.getOccupancyMapWithFutureStatus(occupied_num, cloud_to_publish, &future_status[0][0], 0.6);
//    my_map.getOccupancyMap(occupied_num, cloud_to_publish, 0.6);

    /// Publish Point cloud and center position
    pcl::toROSMsg(cloud_to_publish, cloud_to_pub_transformed);
    cloud_to_pub_transformed.header.frame_id = "map";
    cloud_to_pub_transformed.header.stamp = ros::Time::now();
    cloud_pub.publish(cloud_to_pub_transformed);

    geometry_msgs::PoseStamped map_pose;
    map_pose.header.stamp = cloud_to_pub_transformed.header.stamp;
    map_pose.pose.position.x = uav_position.x();
    map_pose.pose.position.y = uav_position.y();
    map_pose.pose.position.z = uav_position.z();
    map_pose.pose.orientation.x = uav_att.x();
    map_pose.pose.orientation.y = uav_att.y();
    map_pose.pose.orientation.z = uav_att.z();
    map_pose.pose.orientation.w = uav_att.w();
    map_center_pub.publish(map_pose);

    gazebo_model_states_pub.publish(ground_truth_model_states_this);

    /// Publish future status of one layer
    pcl::PointCloud<pcl::PointXYZRGB> future_status_cloud;
    static const int z_index_to_show = MAP_HEIGHT_VOXEL_NUM / 2 - 2; ///Layer
    for(int i=0; i<MAP_LENGTH_VOXEL_NUM; ++i){
        for(int j=0; j<MAP_WIDTH_VOXEL_NUM; ++j){
            int index_this = i*MAP_WIDTH_VOXEL_NUM*MAP_HEIGHT_VOXEL_NUM + j*MAP_HEIGHT_VOXEL_NUM + z_index_to_show;

            for(int n=0; n<PREDICTION_TIMES; ++n){
                pcl::PointXYZRGB p_this;
                float x_offset = (float)n * 12.f; //Used to show prediction at different times in one map

                my_map.getVoxelPositionFromIndexPublic(index_this, p_this.x, p_this.y, p_this.z);
                p_this.x += x_offset;

                float weight_this = future_status[index_this][n];
                int r, g, b;
                colorAssign(r, g, b, weight_this, 0.f, 0.1f, 1);
                p_this.r = r;
                p_this.g = g;
                p_this.b = b;
                future_status_cloud.push_back(p_this);
            }
        }
    }

    sensor_msgs::PointCloud2 cloud_future_transformed;
    pcl::toROSMsg(future_status_cloud, cloud_future_transformed);
    cloud_future_transformed.header.frame_id = "map";
    cloud_future_transformed.header.stamp = ros::Time::now();
    future_status_pub.publish(cloud_future_transformed);

//    future_status_pub

    finish2 = clock();
    double duration2 = (double)(finish2 - start2) / CLOCKS_PER_SEC;
    printf( "****** Map publish time %f seconds\n \n", duration2);

    last_time = this_time;
}

void simObjectStateCallback(const gazebo_msgs::ModelStates &msg)
{
    ground_truth_model_states = msg;
    ground_truth_updated = 1;
}

void simOdomCallback(const nav_msgs::Odometry &msg)
{
    Eigen::Vector3f uav_position_global;
    Eigen::Quaternionf uav_att_global;

    uav_position_global.x() = msg.pose.pose.position.x;
    uav_position_global.y() = msg.pose.pose.position.y;
    uav_position_global.z() = msg.pose.pose.position.z;

    uav_att_global.x() = msg.pose.pose.orientation.x;
    uav_att_global.y() = msg.pose.pose.orientation.y;
    uav_att_global.z() = msg.pose.pose.orientation.z;
    uav_att_global.w() = msg.pose.pose.orientation.w;

    uav_position_global_queue.push(uav_position_global);
    uav_att_global_queue.push(uav_att_global);
    pose_att_time_queue.push(msg.header.stamp.toSec());
}


int main(int argc, char **argv)
{
    ros::init(argc, argv, "planning_test");
    ros::NodeHandle n;

    my_map.setPredictionVariance(0.05, 0.1);  //0.1, 0.1
    my_map.setObservationStdDev(0.01);
//    my_map.setParticleRecordFlag(1, 22.5);
    my_map.setNewBornParticleNumberofEachPoint(30);
    my_map.setNewBornParticleWeight(0.01);
    my_map.setOriginalVoxelFilterResolution(res);

    ros::Subscriber object_states_sub = n.subscribe("/gazebo/model_states", 1, simObjectStateCallback);
    ros::Subscriber point_cloud_sub = n.subscribe("/iris/vi_sensor/camera_depth/depth/points", 1, cloudCallback);
//    ros::Subscriber pose_sub = n.subscribe("/iris/ground_truth/pose", 1, simPoseCallback);
    ros::Subscriber odom_sub = n.subscribe("/iris/ground_truth/odometry", 1, simOdomCallback);

    cloud_pub = n.advertise<sensor_msgs::PointCloud2>("/my_map/cloud_ob", 1, true);
    map_center_pub = n.advertise<geometry_msgs::PoseStamped>("/my_map/map_center", 1, true);
    gazebo_model_states_pub = n.advertise<gazebo_msgs::ModelStates>("/my_map/model_states", 1, true);

    future_status_pub = n.advertise<sensor_msgs::PointCloud2>("/my_map/future_status", 1, true);

    current_velocity_pub = n.advertise<visualization_msgs::MarkerArray>("/my_map/velocity_marker", 1);
    single_object_velocity_pub = n.advertise<geometry_msgs::TwistStamped>("/my_map/single_object_velocity", 1);
    single_object_velocity_truth_pub = n.advertise<geometry_msgs::TwistStamped>("/my_map/single_object_velocity_ground_truth", 1);

//    ros::spin();

    ros::AsyncSpinner spinner(3); // Use 3 threads
    spinner.start();
    ros::waitForShutdown();

    return 0;
}
